import streamlit as st
import requests
import json
import time
import pandas as pd
from datetime import datetime
import uuid
import random
import string

# 设置页面配置
st.set_page_config(
    page_title="数据库问答助手",
    page_icon="🤖",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# 设置API信息
DIFY_API_URL = "http://brainy-agent.sinohealth.cn/v1/chat-messages"
DIFY_API_KEY = "app-F5T16gjHrvgBmNy8jrYb4L4a"
# ByteHouse数据库连接信息
BYTEHOUSE_HOST = 'tenant-2102622389-cn-guangzhou-public.bytehouse.volces.com'
BYTEHOUSE_PORT = 3306
BYTEHOUSE_DATABASE = 'strategy_cloud'
BYTEHOUSE_USERNAME = 'bytehouse'
BYTEHOUSE_PASSWORD = 'gd3gjGTXTb:WBSEjA67dI'

# 设置页面标题
st.title("中康瓴合sql 查询测试 🤖 🔍")
# st.markdown("---")

# 添加示例问题提示
st.markdown("### 👇 您可以尝试以下问题:")
st.markdown("""
- **你好**
- **帮我分析一下**
- **一品红公司过去五年奥司他韦的销售情况如何？**
""")

# 初始化会话状态
if "messages" not in st.session_state:
    st.session_state.messages = []

# 使用UUID生成随机用户ID
def generate_random_user_id():
    # 生成一个UUID，更加唯一和安全
    return f"user-{uuid.uuid4()}"

# 每次运行时重新生成user_id
user_id = generate_random_user_id()

# 侧边栏设置
# with st.sidebar:
#     st.header("设置")
    
#     # 用户ID设置
#     st.subheader("用户信息")
#     user_id = st.text_input("用户ID", user_id)
    
#     # 显示数据库连接信息（只读）
#     st.subheader("数据库连接信息")
#     st.info(f"已连接到 ByteHouse 数据库: {BYTEHOUSE_DATABASE}")
    
#     # 清空对话按钮
#     if st.button("清空对话历史"):
#         st.session_state.messages = []
#         st.success("对话历史已清空！")

# 显示对话历史
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        if message["role"] == "assistant" and "data" in message:
            # 如果有表格数据，显示为DataFrame
            if isinstance(message["data"], list) and len(message["data"]) > 0:
                st.write("📊 查询结果：")
                st.dataframe(pd.DataFrame(message["data"]))
        
        # 显示消息内容
        st.markdown(message["content"])
        
        # 如果是助手消息且有耗时信息，显示耗时
        if message["role"] == "assistant" and "elapsed_time" in message:
            st.caption(f"⏱️ 耗时: {message['elapsed_time']:.2f} 秒")

# 用户输入
prompt = st.chat_input("请输入您的问题...")

# 处理用户输入
if prompt:
    # 添加用户消息到历史
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # 显示用户消息
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # 显示助手正在思考
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        
        # 记录开始时间
        start_time = time.time()
        
        # 准备请求数据
        payload = {
            "inputs": {
                "host": BYTEHOUSE_HOST,
                "port": BYTEHOUSE_PORT,
                "database": BYTEHOUSE_DATABASE,
                "username": BYTEHOUSE_USERNAME,
                "password": BYTEHOUSE_PASSWORD,
                "is_disabled_reporter":1
            },
            "query": prompt,
            "response_mode": "streaming",
            "conversation_id": "",
            "user": user_id
        }
        
        headers = {
            "Authorization": f"Bearer {DIFY_API_KEY}",
            "Content-Type": "application/json"
        }
        
        try:
            with requests.post(
                DIFY_API_URL,
                headers=headers,
                data=json.dumps(payload),
                stream=True
            ) as response:
                
                # 检查响应
                if response.status_code == 200:
                    # 初始化变量
                    full_answer = ""
                    data = None
                    
                    # 处理流式响应
                    for line in response.iter_lines():
                        if line:
                            # 解析SSE格式数据
                            line_text = line.decode('utf-8')
                            if line_text.startswith('data:'):
                                json_str = line_text[5:].strip()
                                if json_str:
                                    try:
                                        chunk = json.loads(json_str)
                                        
                                        # 处理回答片段
                                        if 'answer' in chunk:
                                            answer_chunk = chunk.get('answer', '')
                                            full_answer += answer_chunk
                                            message_placeholder.markdown(full_answer)
                                        
                                        # 处理表格数据
                                        if 'data' in chunk and not data:
                                            try:
                                                data = chunk['data']
                                                if data and isinstance(data, list) and len(data) > 0:
                                                    st.dataframe(pd.DataFrame(data))
                                            except:
                                                pass
                                                
                                        # 处理完成标志
                                        if chunk.get('event') == 'done':
                                            break
                                    except json.JSONDecodeError:
                                        continue
                    
                    # 计算耗时
                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    
                    # 显示耗时
                    st.caption(f"⏱️ 耗时: {elapsed_time:.2f} 秒")
                    
                    # 添加助手消息到历史
                    assistant_message = {"role": "assistant", "content": full_answer, "elapsed_time": elapsed_time}
                    if data:
                        assistant_message["data"] = data
                    st.session_state.messages.append(assistant_message)
                
                else:
                    error_msg = f"请求失败: {response.status_code} - {response.text}"
                    message_placeholder.markdown(f"❌ {error_msg}")
                    st.session_state.messages.append({"role": "assistant", "content": f"❌ {error_msg}"})
                        
        except Exception as e:
            error_msg = f"发生错误: {str(e)}"
            message_placeholder.markdown(f"❌ {error_msg}")
            st.session_state.messages.append({"role": "assistant", "content": f"❌ {error_msg}"})

# 页脚
# st.markdown("---")
# st.markdown(f"© {datetime.now().year} 数据库智能问答系统 | 当前时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
